package com.slipper.ai.flow.support;

import com.slipper.ai.flow.INodeDealer;
import com.slipper.ai.flow.vo.*;
import com.slipper.ai.prompt.Prompts;
import com.slipper.ai.service.ChatService;
import com.slipper.system.ai.bo.AiModelBO;
import com.slipper.system.ai.service.AiModelService;
import com.slipper.util.Tool;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.stream.Collectors;

@Component("aiClassifierNodeDealer")
public class ClassifierNodeDealer implements INodeDealer {
    @Resource
    private AiModelService aiModelService;
    @Resource
    private ChatService chatService;

    @Override
    public String getType() {
        return NodeVO.TYPE_CLASSIFIER;
    }

    @Override
    public List<LineVO> doDeal(FlowVO flow, NodeVO node, NodeInstVO nodeInst) {
        if (!(node instanceof ClassifierNodeVO classifierNode)) throw new RuntimeException("节点类型不正确！");
        if (Tool.isBlank(nodeInst.getInputValue())) throw new RuntimeException("用户问题/输入变量不能为空！");
        if (classifierNode.getClassifierList().isEmpty()) throw new RuntimeException("分类不能为空！");

        String classifierList = classifierNode.getClassifierList()
                .stream()
                .map(ClassifierVO::getContent)
                .filter(Objects::nonNull)
                .flatMap(str -> Arrays.stream(str.split(",")))
                .map(String::trim)
                .filter(str -> !str.isEmpty())
                .collect(Collectors.joining(","));

        Map<String, Object> outputMap = new HashMap<>();
        AiModelBO aiModel = aiModelService.queryBy(classifierNode.getModelId());
        String result = chatService.chat(
                Prompts.CLASSIFIER_PROMPT.formatted(classifierList),
                nodeInst.getInputValue(),
                aiModel);
        int index = -1;
        if (Tool.isBlank(result) || "-1".equals(result)) {
            outputMap.put("分类索引", -1);
            outputMap.put("分类内容", "");
        } else {
            for (int i = 0; i < classifierNode.getClassifierList().size(); i++) {
                String content = classifierNode.getClassifierList().get(i).getContent();
                if (Arrays.stream(content.split(",")).map(String::trim).anyMatch(result::equals)) {
                    index = i;
                    break;
                }
            }
            outputMap.put("分类索引", index + 1);
            outputMap.put("分类内容", result);
        }

        nodeInst.setOutput(outputMap);

        if (index >= 0) {
            ClassifierVO classifierVO = classifierNode.getClassifierList().get(index);
            String id = "source_handle_" + classifierVO.getId();
            return flow.getLines()
                    .stream()
                    .filter(line -> line.getSourceHandle().equals(id))
                    .collect(Collectors.toList());
        } else {
            String id = "source_handle_" + node.getId();
            return flow.getLines()
                    .stream()
                    .filter(line -> line.getSourceHandle().equals(id))
                    .collect(Collectors.toList());
        }
    }
}
